import matplotlib.pyplot as plt
import torch

# -----  生成的曲线 ------------
x = torch.linspace(0, 1, 20)
y = torch.sin(2 * x)  # 0 - 1
plt.plot(x, y, 'ro')


# ------ 激活函数 -----------------
def sigmoid(x):  # 0 - 1
    return 1 / (1 + torch.exp(-x))


def tanh(x):  # -1 - 1
    return (torch.exp(x) - torch.exp(-x)) / (torch.exp(x) + torch.exp(-x))


# ------ 预测线 -------------------
w = torch.tensor(0.2, requires_grad=True)
b = torch.tensor(0.1, requires_grad=True)
predict_y = sigmoid(w * x + b)
line, = plt.plot(x.detach().numpy(), predict_y.detach().numpy(), 'b--')
# ------ 将预测线动态显示-----------
for i in range(1000):
    # 预测
    predict_y = tanh(w * x + b)
    # 更新图线
    line.set_data(x.detach().numpy(), predict_y.detach().numpy())
    # 计算损失函数（横向两线的差距）
    loss = torch.mean((y - predict_y) ** 2)
    loss.backward()  # 反向传播
    # --------- 斜率x小固定值， SGD梯度下降法（随机梯度下降法）
    with torch.no_grad():
        # 获取w、b的斜率
        slope_w = w.grad
        slope_b = b.grad
        # 梯度下降（w,b变化）
        w -= slope_w * 0.5
        b -= slope_b * 0.5
        # 将导数缓存清空
        w.grad.zero_()
        b.grad.zero_()
    plt.pause(0.5)
    print(f"loss:{loss:.4f}")  # 保留4位小数
